pacman::p_load(tidyverse,data.table,broom,hrbrthemes,plyr,latex2exp,openxlsx,ggpubr,ggpmisc)
rm(list=ls())
################################################################################

###Spatial units
spatial_id_a = 'county_id'
spatial_id_b = 'amc_id'
df = fread(paste0('../../data/landuse/clean/geographicunits_argentina.csv.gz'))
setnames(df, old=spatial_id_a, new='spatial_id')
df_a = unique(df[,list(spatial_id,state_id)], by = c("spatial_id","state_id"))
df=fread(paste0('../../data/landuse/clean/geographicunits_brazil.csv.gz'))
setnames(df, old=spatial_id_b, new='spatial_id')
df_b = unique(df[,list(spatial_id,state_id)], by = c("spatial_id","state_id"))
df_ab= unique(rbind(df_a,df_b))
df = setDT(read.xlsx('aux_params/parameters.xlsx',  sheet = "spatial_units"))
df_u_cw = merge(df,df_ab,by='spatial_id', all.x=T, sort=F)[,list(spatial_id_name,state_id)]
df_weights = setDT(read.xlsx('aux_params/parameters.xlsx',  sheet = "weights"))[,list(spatial_id_name,weight_beef)]

###Emissions intensities
df = setDT(read.xlsx('aux_params/parameters.xlsx',  sheet = "emissions_total"))
df = df[, list(spatial_id_name,beef_total_flow_per_t)]
colnames(df)<-c('spatial_id_name','E') 
df_emissions=df

###Baseline model
for(cf_spec in c('dt','ur')){

  for(conduct in c('ic','pc')){
    
    df = setDT(read.csv(file =paste0('results_',conduct,'/lf/Q.csv'),header=FALSE))[,list(V1)]
    setnames(df, old=c('V1'), new=c('Q'))
    df$spatial_id_name = df_emissions$spatial_id_name
    df_lf = df
    
    df = setDT(read.csv(file = paste0('results_',conduct,'/',cf_spec,'/Q.csv'), header=FALSE))[,list(V1)]
    setnames(df, old=c('V1'), new=c('Q'))
    df$spatial_id_name = df_emissions$spatial_id_name
    df_cf = df
    
    df = merge(df_cf, df_lf, by='spatial_id_name', suffixes = c('_cf',''), sort=F)
    df = merge(df, df_emissions, by='spatial_id_name', sort=F)
    df = merge(df, df_weights, by='spatial_id_name', sort=F)

    if(conduct=='ic'){df$conduct_spec='imperfect competition'; df_ic = df}
    if(conduct=='pc'){df$conduct_spec='perfect competition'; df_pc = df}
    
  }
  
  df = rbind(df_ic,df_pc)
  df$x_var=df$E
  df$y_var=100*(df$Q_cf - df$Q)/df$Q
  ylabel=TeX('$\\Delta$ production (pp)')
  df$size_var=df$weight_beef
  df = merge(df, df_u_cw, by='spatial_id_name', all.x=T, sort=F)
  df = df[,list(state_id,conduct_spec,x_var,y_var,size_var)][, lapply(.SD, mean, na.rm=TRUE), by=list(state_id,conduct_spec)]
  
  if(cf_spec=='dt'){plot_title_cf='A. Downstream tax (untargeted)'}
  if(cf_spec=='ur'){plot_title_cf='B. Upstream policy (targeted)'}
  figure = ggplot(data = df, aes(x = x_var, y=y_var, color=conduct_spec) ) +
    geom_point(alpha = 0.25, aes(size=size_var, shape = conduct_spec)) +
    stat_poly_line(formula = y ~ x, aes(weight = size_var, linetype = conduct_spec), se=F, linewidth=1)+
    theme_minimal() +
    labs(title=plot_title_cf, x=TeX('emissions intensity'), y=ylabel)+
    theme(axis.title.x = element_text(hjust = 0.95, size=8), 
          axis.title.y = element_text(hjust = 0.9, size=8),
          axis.text.x = element_text(size=8),
          axis.text.y = element_text(size=8),
          plot.title = element_text(size = 9, hjust = 0.5, vjust=1),
          aspect.ratio = 1,
          panel.grid.minor = element_blank(),
          legend.title = element_blank(),
          legend.key.width = unit(1.5,"cm"),
          legend.position = 'bottom',
          legend.text=element_text(size=8))+
    guides(size = "none", shape = guide_legend(override.aes = list(size = 4.5)))+
    scale_color_manual(values=c('darkred','gray5'))+
    scale_linetype_manual(values=c('solid','dotted'))
  
  figure_bw <- figure + scale_color_grey(start = 0, end = 0.3)
  
  if(cf_spec=='dt'){fig_dt=figure;fig_dt_bw=figure_bw }
  if(cf_spec=='ur'){fig_ur=figure;fig_ur_bw=figure_bw}
  
}
figure4 =  ggarrange(fig_dt,fig_ur, nrow=1, ncol=2, common.legend = TRUE, legend='bottom')
ggsave(paste0("../../output/figures/figure4.pdf"), width = 5.5, height = 3, units='in')
figure4_bw =  ggarrange(fig_dt_bw,fig_ur_bw, nrow=1, ncol=2, common.legend = TRUE, legend='bottom')
ggsave(paste0("../../output/figures/figure4_bw.pdf"), width = 5.5, height = 3, units='in')


###Mobility extension
for(mobility_spec in c('baseline','mobility')){
  
  if(mobility_spec=='baseline'){folder_lf=paste0('results_ic/lf/'); folder_cf=paste0('results_ic/dt/')
  }else{folder_lf=paste0('results_ic/lf/mobility/');folder_cf=paste0('results_ic/dt/mobility/')}
  
  df = setDT(read.csv(file = paste0(folder_lf,'Q.csv'), header=FALSE))[,list(V1)]
  setnames(df, old=c('V1'), new=c('Q') )
  df$spatial_id_name = df_emissions$spatial_id_name
  df_lf = df
  
  df = setDT(read.csv(file = paste0(folder_cf,'Q.csv'), header=FALSE))[,list(V1)]
  setnames(df, old=c('V1'), new=c('Q') )
  df$spatial_id_name = df_emissions$spatial_id_name
  df_cf = df
  
  df = merge(df_cf, df_lf, by='spatial_id_name', suffixes = c('_cf',''), sort=F)
  df = merge(df, df_emissions, by='spatial_id_name', sort=F)
  df = merge(df, df_weights, by='spatial_id_name', sort=F)
  if(mobility_spec=='baseline'){df$mobility_spec='1. no migration (baseline)';df_nm = df}
  if(mobility_spec=='mobility'){df$mobility_spec='2. with migration (extension)';df_wm = df}
}
df = rbind(df_nm,df_wm)
df$x_var=df$E
df$y_var=100*(df$Q_cf - df$Q)/df$Q
ylabel=TeX('$\\Delta$ (pp)')
df$size_var=df$weight_beef
df = merge(df, df_u_cw, by='spatial_id_name', all.x=T, sort=F)
df = df[,list(state_id,mobility_spec,x_var,y_var,size_var)][, lapply(.SD, mean, na.rm=TRUE), by=list(state_id,mobility_spec)]
figure = ggplot(data = df, aes(x = x_var, y=y_var, color=mobility_spec) ) +
  geom_point(alpha = 0.05, aes(size=size_var)) +
  stat_poly_line(formula = y ~ x, aes(weight = size_var, linetype = mobility_spec), se=F, linewidth=1)+
  theme_minimal() +
  labs(title='', x=TeX('emissions intensity'), y=TeX('$\\Delta$ production (pp)') )+
  theme(axis.title.x = element_text(hjust = 0.95, size=8.5), 
        axis.title.y = element_text(hjust = 0.95, size=8.5),
        plot.title = element_blank(),
        axis.text.x = element_text(size=9),
        axis.text.y = element_text(size=9),
        aspect.ratio = 1,
        panel.grid.minor = element_blank(),
        legend.title = element_blank(),
        legend.position = 'bottom',
        legend.text=element_text(size=10))+
  guides(size = "none")+
  scale_color_manual(values=c('darkred','darkred'))+
  scale_linetype_manual(values=c('solid','twodash'))
ggsave(paste0("../../output/figures/appendix/figure8_mobility.pdf"), width = 4.5, height = 4)


